"""
数据源同步任务
"""
import logging
from datetime import datetime, timedelta
from django.utils import timezone
from django.db import transaction
from celery import shared_task
from .utils import AkShareClient
from .models import StockBasic, StockDaily, FinancialIndicator

logger = logging.getLogger(__name__)


@shared_task
def sync_stock_basic():
    """
    同步股票基本信息的Celery任务
    """
    try:
        # 获取股票列表
        stock_df = ak.stock_basic_em()
        count = 0
        
        for _, row in stock_df.iterrows():
            # 使用get_or_create避免重复
            stock, created = StockBasic.objects.update_or_create(
                ts_code=row['ts_code'],
                defaults={
                    'symbol': row['symbol'],
                    'name': row['name'],
                }
            )
            
            if created:
                count += 1
        
        # 获取行业信息
        industry_df = ak.stock_industry_em()
        if not industry_df.empty:
            for _, row in industry_df.iterrows():
                symbol = row.get('代码', '')
                if not symbol:
                    continue
                
                # 转换为ts_code格式
                ts_code = f"{symbol}.SH" if symbol.startswith(('6', '9')) else f"{symbol}.SZ"
                
                try:
                    stock = StockBasic.objects.get(ts_code=ts_code)
                    stock.industry = row.get('所属行业', '')
                    stock.save(update_fields=['industry'])
                except StockBasic.DoesNotExist:
                    pass
        
        return f"成功同步 {count} 只新股票，更新 {len(stock_df) - count} 只股票的基本信息"
    
    except Exception as e:
        logger.error(f"同步股票基本信息失败: {str(e)}")
        raise


@shared_task
def sync_daily_data(ts_code=None, days=30):
    """
    同步股票日线数据的Celery任务
    
    Args:
        ts_code (str, optional): 股票代码，如 '000001.SZ'，为None则同步所有股票
        days (int, optional): 同步最近多少天数据，默认30天
    """
    try:
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
        count = 0
        
        # 获取待同步的股票列表
        if ts_code:
            stocks = StockBasic.objects.filter(ts_code=ts_code)
        else:
            stocks = StockBasic.objects.all()
        
        for stock in stocks:
            try:
                # 获取日线数据
                symbol = stock.symbol
                daily_df = ak.stock_zh_a_hist(symbol, start_date, end_date)
                
                if daily_df.empty:
                    continue
                
                # 批量创建对象
                daily_objects = []
                for _, row in daily_df.iterrows():
                    daily_objects.append(
                        StockDaily(
                            stock=stock,
                            trade_date=row['trade_date'],
                            open=row['open'],
                            high=row['high'],
                            low=row['low'],
                            close=row['close'],
                            change=row.get('change', 0),
                            pct_chg=row.get('pct_chg', 0),
                            vol=row.get('vol', 0),
                            amount=row.get('amount', 0),
                        )
                    )
                
                # 批量更新或创建
                for obj in daily_objects:
                    StockDaily.objects.update_or_create(
                        stock=obj.stock,
                        trade_date=obj.trade_date,
                        defaults={
                            'open': obj.open,
                            'high': obj.high,
                            'low': obj.low,
                            'close': obj.close,
                            'change': obj.change,
                            'pct_chg': obj.pct_chg,
                            'vol': obj.vol,
                            'amount': obj.amount,
                        }
                    )
                
                count += 1
                logger.info(f"成功同步股票 {stock.ts_code} 的日线数据")
            
            except Exception as e:
                logger.error(f"同步股票 {stock.ts_code} 日线数据失败: {str(e)}")
                continue
        
        return f"成功同步 {count} 只股票的日线数据"
    
    except Exception as e:
        logger.error(f"同步日线数据任务失败: {str(e)}")
        raise


@shared_task
def sync_financial_data(ts_code=None, period=20231231):
    """
    同步财务数据的Celery任务
    
    Args:
        ts_code (str, optional): 股票代码，如 '000001.SZ'，为None则同步所有股票
        period (int, optional): 报告期，默认2023年年报
    """
    try:
        # 获取待同步的股票列表
        if ts_code:
            stocks = StockBasic.objects.filter(ts_code=ts_code)
        else:
            stocks = StockBasic.objects.all()
        
        count = 0
        for stock in stocks:
            try:
                # 获取财务指标数据
                financial_df = ak.stock_financial_indicator_em(stock.ts_code, period)
                
                if financial_df.empty:
                    continue
                
                # TODO: 处理财务数据，保存到数据库
                # 具体逻辑取决于AkShare返回的数据结构
                
                count += 1
                logger.info(f"成功同步股票 {stock.ts_code} 的财务数据")
            
            except Exception as e:
                logger.error(f"同步股票 {stock.ts_code} 财务数据失败: {str(e)}")
                continue
        
        return f"成功同步 {count} 只股票的财务数据"
    
    except Exception as e:
        logger.error(f"同步财务数据任务失败: {str(e)}")
        raise 